import os
import numpy as np
import sys
from torch.utils.data import DataLoader
from tqdm import tqdm
import torch
import torch.nn as nn
import torch.optim as optim
from data_loader import iris_dataloader


#初始化神经网络模型
class NN(nn.Module):
    def __init__(self,in_dim,hidden_dim1,hidden_dim2,out_dim):
        # 初始化函数，定义网络结构
        super().__init__()
        # 调用父类的初始化函数
        self.layer1 = nn.Linear(in_dim,hidden_dim1)
        # 定义第一层线性层，输入维度为in_dim，输出维度为hidden_dim1
        self.layer2 = nn.Linear(hidden_dim1,hidden_dim2)
        # 定义第二层线性层，输入维度为hidden_dim1，输出维度为hidden_dim2
        self.layer3 = nn.Linear(hidden_dim2,out_dim)
    def forward(self,x):
        # 将输入x传入第一层
        x = self.layer1(x)
        # 将第一层的输出传入第二层
        x = self.layer2(x)
        # 将第二层的输出传入第三层
        x = self.layer3(x)
        # 返回第三层的输出
        return x
#训练神经网络模型
device = torch.device('mps' if torch.backends.mps.is_available() else 'cpu') # 如果MPS可用，则使用MPS设备，否则使用CPU设备
print(device)
#训练集、验证集、测试集
custom_dataset = iris_dataloader(r"/Users/xiaoyong/PycharmProjects/pytorch-learning/鸢尾花/Iris_data.txt")
# 将数据集划分为训练集、验证集和测试集
train_size = int(len(custom_dataset) * 0.5 )
val_size = int(len(custom_dataset) * 0.3 )
test_size = len(custom_dataset) - train_size - val_size
train_dataset, val_dataset, test_dataset = torch.utils.data.random_split(custom_dataset, [train_size, val_size, test_size])
# 创建数据加载器
train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=1, shuffle=False)
test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)

print(f'训练集的大小：{len(train_loader) * 16}, 验证集的大小：{len(val_loader)}, 测试集的大小：{len(test_loader)}')

#定义一个推理函数，来计算并返回准确率
def infer(model, dataset, device):
    # 将模型设置为评估模式
    model.eval()
    # 初始化正确预测的样本数量
    acc_num = 0
    # 不计算梯度
    with torch.no_grad():
        # 遍历数据加载器中的数据
        for data in dataset:
            datas,label = data
            outputs = model(datas.to(device))
            predict_y = torch.max(outputs, dim=1)[1]
            acc_num += torch.eq(predict_y, label.to(device)).sum().item()
    # 计算准确率
    acc = acc_num / len(dataset)
    return acc
#定义训练函数
def main(lr=0.005,epochs=20):
    # 定义模型
    model = NN(4,12,6,3).to(device)
    # 定义损失函数
    loss_f = nn.CrossEntropyLoss()
    pg = [p for p in model.parameters() if p.requires_grad] #获取需要优化的参数
    # 定义优化器
    optimizer = optim.Adam(pg, lr=lr)
    #权重文件存储路径
    save_path = os.path.join(os.getcwd(),'result/weights')
    if os.path.exists(save_path) is False:
        os.makedirs(save_path)
    #训练过程
    for epoch in range(epochs):
        model.train()#将模型设置为训练模式
        acc_num = torch.zeros(1).to(device)#初始化正确预测的样本数量
        sample_num = 0#初始化样本数量
        train_bar = tqdm(train_loader,file=sys.stdout,ncols=100) #创建进度条
        for datas in train_bar:
            data,label = datas #获取数据和标签
            label = label.squeeze(-1) #将标签从二维张量转换为1维张量
            sample_num += data.shape[0] #更新样本数量

            optimizer.zero_grad()
            outputs = model(data.to(device))
            pred_class = torch.max(outputs,dim=1)[1] #获取预测类别，torch.max的返回值为一个元组，第一个值为Max的值，第二个元素为值的索引
            acc_num += torch.eq(pred_class,label.to(device)).sum().item()

            loss = loss_f(outputs,label.to(device))
            loss.backward()
            optimizer.step()

            train_acc = acc_num / sample_num
            train_bar.desc = "train epoch [{}/{}) loss:{:.3f}".format(epoch+1,epochs,loss)
        #验证过程
        val_acc = infer(model,val_loader,device)
        print('[epoch %d] train_loss: %.3f  train_acc: %.3f  val_accuracy: %.3f' %  (epoch + 1, loss, train_acc, val_acc))
        #保存模型
        torch.save(model.state_dict(),os.path.join(save_path,'nn.pth'))

        #每次数据集跌代之后要对初始化的指标清零
        train_acc = 0
        val_acc = 0
    print('训练完成')
    #测试过程
    test_acc = infer(model,test_loader,device)
    print("test_acc:",test_acc)

if __name__ == '__main__':
    main()

